import networkx as nx
import time

import numpy as np
import math
import random

from projectq.backends import CommandPrinter,Simulator
from projectq.cengines import (MainEngine,
                               DecompositionRuleSet)
import projectq.setups.decompositions
from projectq.setups import restrictedgateset

from projectq.meta import Control
from projectq.ops import H,Measure,All,Rz,Ry,CNOT,Allocate,CX,Deallocate


#####无向图相关#####
class MapGraph(object):

    def __init__(self, mapper):
        self.Graph, self.MaxQuregLength = self.undirected_graph(mapper)
        

    def undirected_graph(self, mapper):
        G = nx.path_graph(0)
        lines = mapper.split('\n')
        n = int(lines[0])
        lines = lines[1:]
        for line in lines:
            a, b = line.split(',')
            nx.add_path(G, [int(a)-1, int(b)-1])
        return G, n

    def find_shortest_path(self, sour, targ):
        list_path = nx.shortest_path(self.Graph, source=sour, target=targ)
        # print('%d节点到%d节点最短路径: ' % (sour, targ), list_path)
        return list_path

    def TransformCnot(self, c, t):
        path = self.find_shortest_path(c, t)
        n = len(path)
        if(n==2):
            return [[c, t]]
        ctPairList = []
        cnotList = []

        for i in range(n-1):
            ctPairList.append([path[i],path[i+1]])
        cnotList += ctPairList
        cnotList += ctPairList[-2::-1]
        cnotList += ctPairList[1:]
        cnotList += ctPairList[-2:0:-1]
        return cnotList

#####另一种求uniformly control gate的分解线路图中的Rz或Ry角度的相关方法#####
#####未使用，在态的维度大于2^15之后，求旋转角的变换矩阵会非常耗内存##########
def GetTransformMat(k):
    d = 1<<k
    M = np.ones((d, d), dtype='int')

    j = np.array(range(d), dtype='int')
    for i in range(d):
        # Decimal to Gray Code: x^(x>>1)
        g = i^(i>>1)
        n = j&g
        count = np.zeros(d, dtype='int')
        while(n[n>0].shape[0]):
            mask = (n>0)
            count[mask] = count[mask] + 1
            n = n&(n-1)
            
        mask = (count&1>0)
        M[i,mask] = -1
    return M

def Alpha2Theta(alpha, M):
    alpha = np.array(alpha)
    theta = M.dot(alpha)/len(alpha)
    return -theta

def GetAlphaZ(n, k, phaseAngle):
    alpha = np.zeros((1<<(n-k), 1))
    d = 1<<(k-1)

    l = np.array(range(1, d+1))
    for j in range(1, len(alpha)+1):
        index1 = (2*j-1)*d+l-1
        index2 = index1 - d
        temp = phaseAngle[index1] - phaseAngle[index2]
        alpha[j-1] = temp.sum()
        
    return alpha/d

def GetAlphaY(n, k, realAmpl):
    d = 1<<k
    alpha = np.zeros((1<<(n-k), 1))
    l1 = np.array(range(1,d+1))
    l2 = np.array(range(1,(d>>1)+1))
    for j in range(1, len(alpha)+1):
        index1 = (j-1)*d+l1-1
        temp1 = np.power(realAmpl[index1], 2)
        temp1 = temp1.sum()
        if(temp1 > 1e-12):
            index2 = (2*j-1)*(d>>1)+l2-1
            temp2 = np.power(realAmpl[index2], 2)
            temp2 = temp2.sum()
            alpha[j-1] = 2*math.asin(math.sqrt(temp2/temp1))
    
    return alpha

#####ProjectQ中实现的用递归求解Rz以及Ry门角度的方法#####
def GetRotationTheta(nQubits, targetState, RGate):
    theta = []
    thetaAll = []
    globalPhase = 0
    def GetTheta(angles, nControlQubits, rightmost_cnot):
        if nControlQubits == 0:
            theta.append(angles[0])
        else:
            if rightmost_cnot[nControlQubits]:
                angles1 = []
                angles2 = []
                for lower_bits in range(2**(nControlQubits-1)):
                    leading_0 = angles[lower_bits]
                    leading_1 = angles[lower_bits + 2**(nControlQubits-1)]
                    angles1.append((leading_0 + leading_1)/2.)
                    angles2.append((leading_0 - leading_1)/2.)
            else:
                angles1 = []
                angles2 = []
                for lower_bits in range(2**(nControlQubits-1)):
                    leading_0 = angles[lower_bits]
                    leading_1 = angles[lower_bits + 2**(nControlQubits-1)]
                    angles1.append((leading_0 - leading_1)/2.)
                    angles2.append((leading_0 + leading_1)/2.)

            GetTheta(angles1, nControlQubits-1, rightmost_cnot)
            GetTheta(angles2, nControlQubits-1, rightmost_cnot)
            # Next iteration on this level do the other cnot placement
            rightmost_cnot[nControlQubits] = (not rightmost_cnot[nControlQubits])

    nControlQubits = nQubits
    phase_of_blocks = np.angle(targetState)
    abs_of_blocks = np.real(np.array(targetState) * np.exp(-phase_of_blocks * 1j))
    if(RGate == Rz):
        for target_qubit in range(nQubits):
            angles = []
            phase_of_next_blocks = []
            for block in range(2 ** (nQubits - target_qubit - 1)):
                phase0 = phase_of_blocks[2 * block]
                phase1 = phase_of_blocks[2 * block + 1]
                angles.append(phase0 - phase1)
                phase_of_next_blocks.append((phase0 + phase1) / 2.)
            phase_of_blocks = phase_of_next_blocks

            nControlQubits = nControlQubits - 1
            if(nControlQubits != 0):
                angles1 = []
                angles2 = []
                for lower_bits in range(2 ** (nControlQubits - 1)):
                    leading_0 = angles[lower_bits]
                    leading_1 = angles[lower_bits + 2 ** (nControlQubits - 1)]
                    angles1.append((leading_0 + leading_1) / 2.)
                    angles2.append((leading_0 - leading_1) / 2.)
                rightmost_cnot = {}
                for i in range(nControlQubits + 1):
                    rightmost_cnot[i] = True
                GetTheta(angles1, nControlQubits - 1, rightmost_cnot)
                GetTheta(angles2, nControlQubits - 1, rightmost_cnot)
            else:
                theta = angles

            thetaAll.append(theta)
            theta = []
        globalPhase = phase_of_blocks[0]
    if(RGate == Ry):
        for target_qubit in range(nQubits):
            angles = []
            abs_of_next_blocks = []
            for block in range(2 ** (nQubits - target_qubit - 1)):
                a0 = abs_of_blocks[2 * block]
                a1 = abs_of_blocks[2 * block + 1]
                if np.abs(a0) <1e-10 and np.abs(a1) < 1e-10:
                    angles.append(0)
                else:
                    angles.append(
                        -2. * math.acos(a0 / math.sqrt(a0 ** 2 + a1 ** 2)))
                abs_of_next_blocks.append(math.sqrt(a0 ** 2 + a1 ** 2))
            abs_of_blocks = abs_of_next_blocks

            nControlQubits = nControlQubits - 1
            if (nControlQubits != 0):
                angles1 = []
                angles2 = []
                for lower_bits in range(2 ** (nControlQubits - 1)):
                    leading_0 = angles[lower_bits]
                    leading_1 = angles[lower_bits + 2 ** (nControlQubits - 1)]
                    angles1.append((leading_0 + leading_1) / 2.)
                    angles2.append((leading_0 - leading_1) / 2.)
                rightmost_cnot = {}
                for i in range(nControlQubits + 1):
                    rightmost_cnot[i] = True
                GetTheta(angles1, nControlQubits - 1, rightmost_cnot)
                GetTheta(angles2, nControlQubits - 1, rightmost_cnot)
            else:
                theta = angles

            thetaAll.append(theta)
            theta = []

    return thetaAll, globalPhase

#####缓存uniformly control gate的分解线路的门操作#####

def GetCnotList(k):
    conVec = np.zeros(2**k, dtype=int)
    conVec[0] = 1
    conVec[1] = 1
    for i in range(1,k):
        conVec[2**i-1] += 1
        conVec[2**i:2**(i+1)] = conVec[0:2**i] 
    return conVec

def UniformControlGateDagger(mapGraph, n, k, theta, RGate):
    gateList = []
    if(RGate == Ry):
        strRGate = 'Ry'
    if(RGate == Rz):
        strRGate = 'Rz'

    if k==0:
        if(abs(theta[0]) > 1e-6 ):
            angle = round(float(-theta[0]) % (4. * math.pi), 12)
            # RGate(theta[0]).get_inverse() | quReg[n-1-k]
            return [r"%s(%.6f) | qubit[%d],"%(strRGate, angle, n-1-k)]
    else:
        CnotList = GetCnotList(k)
        for i in range(2**k):
            if(abs(theta[i]) > 1e-6 ):
                angle = round(float(-theta[i]) % (4. * math.pi), 12)
                gateList.append( r"%s(%.6f) | qubit[%d],"%(strRGate, angle, n-1-k) )
                # RGate(theta[i]).get_inverse() | quReg[n-1-k]

            mapCnotList = mapGraph.TransformCnot(CnotList[i]+n-1-k, n-1-k)
            for ct in mapCnotList:
                gateList.append( r"CNOT | (qubit[%d], qubit[%d]),"%(ct[0], ct[1]) )
                # CNOT | (quReg[ct[0]], quReg[ct[1]])

    return gateList

############

def circuit_generator(eng, target_state, mapper):
    '''
        The function you need to modify.
     Args:
       eng:
         The engine type should be defined in "__main__" function.
       target_state(list):
         A target quantum state vector like: [0.5. 0.5. 0.5. 0.5]. And this
         simple example can be prepared from 'H | qubit[0]; H | qubit[1]'
         where the 'qubit' are in a |0> state at very first.
       mapper(string):
         The mapper of a real quantum chip like: '3\n1,2\n2,3'. This simple
         example of mapper means that there are 3 qubits in this given chip
         and you can play 'CNOT' on 'qubit1' and 'qubit2'. You can also play
         'CNOT' on 'qubit2' and 'qubit3'. But you can not play 'CNOT' on
         'qubit1' and 'qubit3' because '1,3' is not in the given mapper.
     Returns:
       simulated_circuit(string):
         After your calculation you may get a final quantum circuit like:
         'H | qubit[0]; H | qubit[1]'. Each step seperated by a '; '. In
         our score we will search operations and qubit index in your result
         and run it on our backend. We will compare the final state your
         circuit produces with our target state by ourselves.
    '''

    gateList = []

    phaseAngle = np.angle(target_state)
    # realAmpl = np.real(np.array(target_state)*np.exp(-phaseAngle*1j))
    n = int(math.log2(len(target_state)))

    #mapper无向图定义
    G = MapGraph(mapper)

    # transMat = GetTransformMat(n-1)

    # for k in range(n):
    #     alphaY = GetAlphaY(n, n-k, realAmpl)
    #     thetaY = Alpha2Theta(alphaY, transMat[:len(alphaY),:len(alphaY)])
    #     gates = UniformControlGateDagger(G, n, k, thetaY, Ry)
    #     gateList = gateList + gates
    # for k in range(n):
    #     alphaZ = GetAlphaZ(n, n-k, phaseAngle)
    #     thetaZ = Alpha2Theta(alphaZ, transMat[:len(alphaZ),:len(alphaZ)])
    #     gates = UniformControlGateDagger(G, n, k, thetaZ, Rz)
    #     gateList = gateList + gates


    # Cancel global phase
    # Ph(globalPhase) | qureg[-1]
    globalPhase = phaseAngle.sum()/2**n
    angle = round(float(-2*globalPhase) % (4. * math.pi), 12)
    # Rz(-2*globalPhase) | quReg[0]
    gateList.append(r"Rz(%.6f) | qubit[%d],"%(angle, 0))

    thetaAll,_ = GetRotationTheta(n, target_state, Ry)
    for k in range(n):
        thetaY = thetaAll[n-k-1]
        gates = UniformControlGateDagger(G, n, k, thetaY, Ry)
        gateList = gateList + gates

    thetaAll, globalPhase = GetRotationTheta(n, target_state, Rz)

    for k in range(n):
        thetaZ = thetaAll[n-k-1]
        gates = UniformControlGateDagger(G, n, k, thetaZ, Rz)
        gateList = gateList + gates

    simulated_circuit = str(gateList)
    simulated_circuit = simulated_circuit[1:-3]
    simulated_circuit = simulated_circuit.replace('\'', '')
    simulated_circuit = simulated_circuit.replace(',,',';')
    return simulated_circuit

##############


if __name__ == "__main__":

    eng = MainEngine(CommandPrinter())

    mapper = '5\n1,2\n1,3\n2,3\n2,4\n2,5\n3,4\n3,5\n4,5'
    # mapper = '20'
    # for i in range(19):
    #     mapper += '\n%d,%d'%(i+1,i+2)
    # with open('map.txt','r') as f:
    #     mapper = f.read()
    
    state = np.array(range(1<<5),dtype='complex128')
    state = state*np.exp(state*1j)
    state = state/np.linalg.norm(state)
    t0 = time.time()
    quantumGates = circuit_generator(eng, state, mapper)
    print(quantumGates)